import torch

input_size = 4
hidden_size = 4
batch_size = 1
"""
问题：将字符串hello输入到rnn循环神经网络，输出ohlol
"""
#构建输入输出字典
idx2char_1 = ['e', 'h', 'l', 'o']
idx2char_2 = ['e', 'h', 'l', 'o']
#hello
x_data = [1, 0, 2, 2, 3]
#ohlol
y_data = [3, 1, 2, 3, 2]
#y_data = [2, 0, 1, 2, 1]

one_hot_lookup = [[1, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]]
#构造独热向量，此时向量维度为(SeqLen*InputSize) seqlen为单词长度 hello单词长度为5
#对于x_data列表中的每一个元素，使用one_hot_lookup[x]构造独热向量 独热向量:只有一个元素为1其余为0
x_one_hot = [one_hot_lookup[x] for x in x_data]
#view(-1……)保留原始SeqLen，并添加batch_size,input_size两个维度 -1自动计算seqlen=5 batch_size为1
inputs = torch.Tensor(x_one_hot).view(-1, batch_size, input_size)
#将labels转换为（SeqLen*1）的维度 -1自动计算seqlen=5
labels = torch.LongTensor(y_data).view(-1, 1)

class Model(torch.nn.Module):
    def __init__(self, input_size, hidden_size, batch_size):
        super(Model, self).__init__()
        self.batch_size = batch_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        #循环神经网络
        self.rnncell = torch.nn.RNNCell(input_size = self.input_size,
                                        hidden_size = self.hidden_size)

    def forward(self, input, hidden):
        # RNNCell input = (batchsize*inputsize)
        # RNNCell hidden = (batchsize*hiddensize)
        # Ht = rnn(Ht-1,xt)
        hidden = self.rnncell(input, hidden)
        return hidden

    #初始化零向量作为h0，只有此处用到batch_size
    def init_hidden(self):
        return torch.zeros(self.batch_size, self.hidden_size)

net = Model(input_size, hidden_size, batch_size)

#求交叉熵
criterion = torch.nn.CrossEntropyLoss()
#Adam一种改进的随机梯度下降算法
optimizer = torch.optim.Adam(net.parameters(), lr=0.1)

for epoch in range(15):
    #损失及梯度置0，创建前置条件h0
    loss = 0
    optimizer.zero_grad()
    hidden = net.init_hidden()

    print("Predicted string: ",end="")
    #inputs=（seqLen*batchsize*input_size）即（5*1*4）labels = (seqLen*1)即（5*1）
    #每次输入一个四维向量即独热向量代表hello中的一个字母 输出一个四维数字即一个字母到另外一个字母的概率
    #input是按序列取的inputs元素（batchsize*inputsize）
    #label是按序列去的labels元素（1）
    #循环5次
    for input, label in zip(inputs, labels):
        hidden = net(input, hidden)
        #序列的每一项损失都需要累加
        loss += criterion(hidden, label)
        #多分类取最大 idx为最大的下标 hidden表示ehlo中每个字母对应的概率
        _, idx = hidden.max(dim=1)
        print(idx2char_2[idx.item()], end='')

    loss.backward()
    optimizer.step()

    print(", Epoch [%d/15] loss = %.4f" % (epoch+1, loss.item()))